from flask import Flask,request,jsonify
import flask
import mmap
import os,sys
import cv2
import json
import time
import copy
import numpy as np
from types import SimpleNamespace
import pprint
LOCAL = False

CAMS = ['卡扣',
        '插线']

ALL_CAMS = ['卡扣',
            '插线',
            '焊点',
            '间隙',
            '弯管']

# ================ ================
idx_tp = 0
yolo_instance = None
yolo_model_kakou = None
yolo_model_line_la = None

yolo_model_line_sm = None
if not LOCAL:
    yolo_dir = 'D:\\ALGO\\media_line2\\yolov8'

    if yolo_dir in sys.path:
        pass
    else:
        sys.path.append(yolo_dir)

    import importlib
    import importlib.util
    yolo_instance = importlib.import_module('ultralytics')

    yolo_model_kakou_weights = 'D:\\ALGO\\media_line2\\models\\line2\\kakou\\best.pt'
    # yolo_model_line_weights = 'D:\\ALGO\\media_line2\\models\\line2\\line\\yolo_v8_media_line2_chaxian.pt'
    yolo_model_line_weights = 'D:\\ALGO\\media_line2\\models\\line2\\line\\best.pt'
    yolo_model_line_la_weights = 'D:\\ALGO\\media_line2\\models\\line2\\line\\best_la.pt'
    yolo_model_line_sm_weights = 'D:\\ALGO\\media_line2\\models\\line2\\line\\best_sm.pt'
    # ================ ================
else:
    yolo_dir = 'D:\\py\\work\\line3\\media_line2\\yolov8'
    if yolo_dir in sys.path:
        pass
    else:
        sys.path.append(yolo_dir)

    import importlib
    import importlib.util
    yolo_instance = importlib.import_module('ultralytics')

    # ================ ================

    yolo_model_kakou_weights = 'D:\\py\\work\\yolov8\\runs\\detect\\train8\\weights\\best.pt'
    yolo_model_line_weights = 'D:\\py\\work\\yolov8\\runs\\detect\\train10\\weights\\best.pt'


def get_model(model):
    model = yolo_instance.YOLO(model)
    
    def infer(im):
        results = model(im,iou=0.5,conf=0.3)
        return results[0].boxes.data.cpu().numpy()
    return infer
    

    

app = Flask(__name__)

@app.route("/")
def hello_world():
    return "<p>This is yolo algo for media line2.</p>"


# 获取内存映射图片
def get_img(name, file_size):
    m = mmap.mmap(-1, file_size, tagname=name, access=mmap.ACCESS_READ)
    # [0]:默认-1  [1] 文件大小 [2]: id就是文件的标签 [3]: 读写模式
    t = time.time()
    while True:
        try:
            
            if time.time() - t > 0.1:  # add 23.12.21 缩短等待时间
                m.close()  # 关闭后继续循环
                return None
            m.seek(0)  # 设置文件的当前位置
            
            image_data = np.frombuffer(m.read(file_size), dtype=np.uint8)  # 将数据转换为NumPy数组
            if image_data.size == 0:
                continue
            image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)  # 直接使用二进制数组解码
            
            if image is None:
                continue
            
            m.close()
            return image  # 这一步相当于跳出循环

        except FileNotFoundError:
            print(f'未找到文件')
            m.close()  # 关闭后继续循环
            return None
        
# 心跳
@app.route('/heartbeat', methods=['GET'])
def heartbeat():
    return 'OK', 200

@app.route("/json", methods=['POST'])
def get_frame():
    global idx_tp, yolo_model_kakou, yolo_model_line_la,yolo_model_line_sm
    #### 算法返给客户端的字典 ####
    # 出参
    json_data = {"code": "00000", 
                 "success": True, 
                 "msg": "算法调用成功", 
                 "data": {
                     "result": {
                         "total_result": "", 
                         "camera_name": "", 
                         "standard_result": [], 
                         "images": []
                         }
                     }
                }
    idx_tp += 1  # 图片处理计数
    t_once_start = time.time()  # 单词处理图片开始时间记录
    now_idx = copy.deepcopy(idx_tp)  # 本张图片处理的序号
    t_s1 = time.time()
    
    # json解析
    receivedJson = json.loads(str(request.data, encoding="utf-8"),
                              object_hook=lambda d: SimpleNamespace(**d))
    print(f'================received json================')
    pprint.pprint(receivedJson)
    
    # 发送json
    sendJson = json.loads(json.dumps(json_data),
                          object_hook=lambda d: SimpleNamespace(**d))
    
    # 相机名称设置错误
    if receivedJson.camera_data.name not in CAMS:
        sendJson.success = False
        sendJson.msg = f'相机名称设置错误，请在 [{CAMS}]中选择名字'
        send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
        print(f'================send json================')
        pprint.pprint(sendJson)
        return send
        
    imgs = receivedJson.camera_data.images
    
    for im in imgs:
        out_result_image = {}
        out_result_image['id'] = im.id  # 内存映射的标签  这一步中传进来的id有相同的
        out_result_image['name'] = im.name  # 内存映射共享图片名字
        out_result_image['size'] = im.size  # 共享内存的图像大小
        out_result_image['coordinate'] = []
        
        # 解码图像
        t_s2 = time.time()
        ### 在内存映射中捞取相应图片  ###
        image = get_img(im.name, im.size)  # 通过内存名字 和长度 取图 ----> numpy img
        if image is None:
            sendJson.success = False
            sendJson.msg = f'算法无法在共享内存中读取图片'
            send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
            print(f'================send json================')
            pprint.pprint(sendJson)
            return send
        t_len = round(time.time() - t_s2, 3)
        if time.time() - t_once_start > 99999:
            print(f"--@@@@[相机]{receivedJson.camera_data.name} [当前处理图片序号]:{now_idx} [图片名]:{im.name}, >>>>当前检测时间过长,直接跳出")
            break
        
        detect_model = None
        detect_model_la = None
        detect_model_sm = None
        print('======== 相机 ========')
        if image is not None:
            
            if receivedJson.camera_data.name == "卡扣":  # 进行相机区分
                
                print(f'相机请求{receivedJson.camera_data.name}')
                detect_model = yolo_model_kakou
            elif receivedJson.camera_data.name == "插线":
                
                print(f'相机请求{receivedJson.camera_data.name}')
                detect_model_la = yolo_model_line_la
                detect_model_sm = yolo_model_line_sm

            if receivedJson.camera_data.name == "卡扣" and detect_model is None:
                sendJson.success = False
                sendJson.msg = f'算法无法找到检测模型，请检查相机名称设置或联系算法检查模型文件 卡扣'
                send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
                print(f'================send json================')
                pprint.pprint(sendJson)
                return send

            elif receivedJson.camera_data.name == "CAM02" and (detect_model_sm is None or detect_model_la is None):
                sendJson.success = False
                sendJson.msg = f'算法无法找到检测模型，请检查相机名称设置或联系算法检查模型文件 插线'
                send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
                print(f'================send json================')
                pprint.pprint(sendJson)
                return send
            else:
                pass
            # yolo模型
            t_s3 = time.time()

            # 双次检测逻辑
            result = None
            if receivedJson.camera_data.name == "卡扣":
                result = detect_model(image)
                print(f'>>>>>>>>>cam01 result')
                pprint.pprint(result)
            elif receivedJson.camera_data.name == "插线":
                la_result = detect_model_la(image)  # 确定是哪颗相机对应的模型
                print(f'>>>>>>>>>cam02 la result')
                pprint.pprint(la_result)
                if len(la_result) == 0:
                    result = []
                else:
                    kuang = la_result[0]
                    x1, y1, x2, y2 = map(int, kuang[:4])
                    cropped_img = image[y1:y2, x1:x2]
                    
                    # cv2.imwrite(str(time.time())+'.jpg',cropped_img)
                    sm_result = detect_model_sm(cropped_img)
                    print(f'>>>>>>>>>cam02 sm result')
                    pprint.pprint(sm_result)

                    result = []
                    for det in sm_result:
                        cx1, cy1, cx2, cy2 = det[:4]
                        # 转换为原图坐标
                        fx1, fy1 = cx1 + x1, cy1 + y1
                        fx2, fy2 = cx2 + x1, cy2 + y1
                        conf = det[4]
                        cls = det[5]
                        result.append([fx1, fy1, fx2, fy2, conf, cls])

                    pprint.pprint(result)

            # temp_out.extend(out)  # 检测结果 out---> [rect, cls, confidence]
            t_len = round(time.time() - t_s3, 3)
            print(f'==== ==== ==== ====')

            # 算法结果处理
            label_kakou = {
                'kakou1':0,
                'kakou2':1,
                'kakou3':2,
            }
            label_line = {
                'line1':0,
                'line2':1,
                'line3':2,
            }
            # 计数
            end_kakou = {
                'kakou1': 0,
                'kakou2': 0,
                'kakou3': 0
            }
            end_line = {
                'line1': 0,
                'line2': 0,
                'line3': 0
            }
        
            if len(result) == 0:
                sendJson.success = True
                sendJson.msg = f'算法无法检测到目标'
                sendJson.data.result.total_result = 'NG'
                sendJson.data.result.camera_name = receivedJson.camera_data.name
                send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
                print(f'================send json================')
                pprint.pprint(sendJson)
                return send
            

            for x1, y1, x2, y2, conf, label in result:
                if receivedJson.camera_data.name == "卡扣":
                    end_kakou[list(label_kakou.keys())[list(label_kakou.values()).index(int(label))]] = 1
                    print(f'卡扣标签')
                else:
                    end_line[list(label_line.keys())[list(label_line.values()).index(int(label))]] = 1
                    print(f'接线标签')
                    
        
                out_result_image['coordinate'].append({
                                        "standard_id": getattr(receivedJson.camera_data.standards[0],'standard_id',''),
                                        "label": receivedJson.camera_data.standards[0].label,
                                        "threshold": receivedJson.camera_data.standards[0].threshold,
                                        "scores": float(conf),
                                        "algo_result": "YOU",
                                        "color": "GREEN",
                                        "list_draw_data": [f"检测项: {receivedJson.camera_data.standards[0].label}{int(label)}, 结果: '1'，阈值: {receivedJson.camera_data.standards[0].threshold}, 分值: {float(conf)}"],
                                        "type": "RECTANGLE",
                                        "list_point": [
                                            [int(x1),int(y1)],
                                            [int(x2),int(y1)],
                                            [int(x1),int(y2)],
                                            [int(x2),int(y2)],
                                        ]
                                    })
        
        end_kakou = sum(end_kakou.values())
        end_line = sum(end_line.values())
        print(f'==== ==== ==== ====kakou==== ==== ==== ====')
        print(end_kakou)
        print(f'==== ==== ==== ====jiexian==== ==== ==== ====')
        print(end_line)
        algo_result = '0'
        if receivedJson.camera_data.name == "卡扣":
            algo_result = str(end_kakou)
        elif receivedJson.camera_data.name == "插线":
            algo_result = str(end_line)    
        
        max_scores = max([row[4] for row in result])
        total_result = 'OK' if algo_result == receivedJson.camera_data.standards[0].standard else 'NG'
        
        sendJson.data.result.images.append(
            out_result_image
        )
    
    sendJson.data.result.camera_name = receivedJson.camera_data.name
    sendJson.data.result.total_result = total_result
    sendJson.data.result.standard_result.append({
        'standard_id':getattr(receivedJson.camera_data.standards[0],'standard_id',''),
        'label':receivedJson.camera_data.standards[0].label,
        'type':receivedJson.camera_data.standards[0].type,
        'standard':receivedJson.camera_data.standards[0].standard,
        'threshold':receivedJson.camera_data.standards[0].threshold,
        'scores':float(max_scores),
        'algo_result':algo_result,
        'detect_result':total_result,
        
    })
    
    send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
    print(f'================send json================')
    pprint.pprint(sendJson)
    return send
    
    
def model_init(yolo_model_kakou_weights,yolo_model_line_la_weights,yolo_model_line_sm_weights):
    global yolo_model_kakou, yolo_model_line_la,yolo_model_line_sm
    try:
        yolo_model_kakou = get_model(yolo_model_kakou_weights)
        yolo_model_line_la = get_model(yolo_model_line_la_weights)
        yolo_model_line_sm = get_model(yolo_model_line_sm_weights)
    except Exception as e:
        print('模型加载失败，检查模型路径')
        print(e)

def flask_init():
    host = "127.0.0.1"
    port = 8080
    # app.debug=True
    print(app.url_map)
    app.run(host=host, port=port, threaded=False)
    
if __name__ == '__main__':
    model_init(yolo_model_kakou_weights,
               yolo_model_line_la_weights,
               yolo_model_line_sm_weights,)
    flask_init()